import torch
from torchvision import transforms
from PIL import Image
import os

from model import LeNet


def main():
    transform = transforms.Compose(
        [transforms.ToTensor(),
         transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]
    )

    classes = ('0', '1', '2', '3', '4', '5')

    net = LeNet()
    net.load_state_dict(torch.load('./LeNet.pth'))

    # 按类读取验证集中的图片，将预测结果与类名比较，求取运算精度
    root = './data/verification_data'
    count = 0
    sum = 0
    for i in range(6):
        for dirpath, dirnames, filenames in os.walk(str(root + '/' + str(i))):
            for image in filenames:
                img = Image.open(str(root + '/' + str(i) + '/' + str(image)))
                img = transform(img)
                img = torch.unsqueeze(img, dim=0)
                with torch.no_grad():
                    outputs = net(img)
                    predict = torch.max(outputs, dim=1)[1].numpy()
                if int(classes[int(predict)]) == i:
                    count += 1
            sum += len(filenames)
    print(float(count) / float(sum))


if __name__ == '__main__':
    main()
